# ================== 数据加载 ==================
import pandas as pd
from fontTools.misc.cython import returns

import 数据预处理
import 读取数据

stock_codes = [
    '601116.SH', '600036.SH', '601398.SH', '601328.SH', '601288.SH',
    '600919.SH', '600000.SH', '601988.SH', '000001.SZ', '601229.SH'
]

start_date = '20100101'
end_date = '20240326'

test_stock_codes = '601009.SH'
test_start_date = '20240501'
test_end_date = '20250325'
# test_stock_codes = '512820.SH'
test_df = 读取数据.load_data(test_stock_codes, test_start_date, test_end_date)


dfs = []
for code in stock_codes:
    df = 读取数据.load_data(code, start_date, end_date)
    print(f"已加载 {code}，数据维度：{df.shape}")
    dfs.append(df)

# ================== 特征工程 ==================
feature_dfs = []
feature_test = None
features = None

for df in dfs:
    # 确保数据按日期排序
    df = df.sort_values('trade_date')

    # 特征计算
    feature_df, current_features = 数据预处理.calculate_features(df.copy())


    # 记录特征列（只需第一次）
    if features is None:
        feature_test, current_features = 数据预处理.calculate_features(test_df.copy())
        features = current_features
        print(f"\n使用的特征列：{features}")

    feature_dfs.append(feature_df)

# ================== 标签生成 ==================

feature_test['target'] = (feature_test['close'].shift(-1) > feature_test['close']).astype(int)

# 清理无效数据
clean_df = feature_test.dropna(subset=['target']).copy()
feature_test = clean_df
print(clean_df['target'].value_counts(normalize=True).apply(lambda x: f"{x:.1%}"))


for idx, df in enumerate(feature_dfs):
    # 计算目标变量
    df['target'] = (df['close'].shift(-1) > df['close']).astype(int)

    # 清理无效数据
    clean_df = df.dropna(subset=['target']).copy()
    feature_dfs[idx] = clean_df

    # 显示样本分布
    print(f"\n股票 {stock_codes[idx]} 有效样本：{len(clean_df)}")
    print("标签分布：")
    print(clean_df['target'].value_counts(normalize=True).apply(lambda x: f"{x:.1%}"))

 # ================== 模式选择 ==================
# 可选模式: 'combined'（联合训练） 或 'individual'（独立训练）
TRAIN_MODE = 'combined'
import 模型训练
model = 模型训练.train_model(feature_dfs,features,stock_codes,TRAIN_MODE)

import 预测
pred = 预测.predict(feature_test,model,features)

feature_test['signal'] =pred

# ================== Backtrader回测 ==================
import 回测
回测.back_trade(feature_test)
